import numpy as np
from typing import Dict
import matplotlib.pyplot as plt
import copy
import math
import sys
import pathlib

from IDM import IDM
from Constants import *
# 由Constants.py导入常用常量
from Moveable import Moveable,moveableInitialize,moveableUpdate,lane
from scipy.stats import multivariate_normal
sys.path.append(str(pathlib.Path(__file__).parent.parent))
from Dataloader.src import waymo_training_plot_tools as plot_tool
from Dataloader.src import waymo_scenario_classes as class_tool
from Dataloader.src.waymo_transfer_training_to_tf import *
from Dataloader.src.waymo_training_math_tools import *
# import tensorflow

from driver_model.Moveable import Moveable

"""
通过当前车道和目标车道的lanelink，判断是否处在路口
若在路口，执行IDM，同时检查是否要变道（当目标车道有两个时）
"""

class TurningModel:
    ismax = 100 # ve(s) = ve(ismax) if s > ismax
    MAX_BRAKING = 10 # m/s^2
    DELTA = 4 # Acceleration exponent
    def __init__(self,v0, a, b,sigma,delta=DELTA,s0=2, T=1.5,size=20):
        '''
        :param p: politeness factor
        :param db: change lane incentive penalty(罚函数)
        :param gap_min: max safe distance
        :param bsave: max safe braking deceleration
        :param bias_right: bias (m/s^2) to drive right
        '''
        self.size,self.sigma = size,sigma
        global ismax
        global MAX_BRAKING
        self.v0, self.a, self.b, self.s0,self.T = v0, a, b, s0, T
        self.delta = delta
        self.sqrtab = np.sqrt(a*b)
    
    def isCrossing(self,me):
        # 判断是否处在路口
        if(len(me.lane.crossingId[0])>1):
            return True
        else:
            return False
    
    def turningDecision(self,me:Moveable, 
                    targetLane:Dict,fwd:Dict[int,Moveable]):
        """利用IDM计算未来位置

        Args:
            me (Moveable): _description_
            targetLane (Dict): 路口所有target lane的字典，形式为{laneId:pathDict}(pathDict为{pathX:list,pathY:list})——以后pathDict都用这种形式
            fwd (Moveable): _description_
            fwd2 (Moveable): _description_
        """
        predictPath = {}
        print("me pos:",me.pos)
        print("me vel",me.vel)
        for key,value in targetLane.items():
            # 对路口的每条可能车道都生成一条预测路径
            idm = IDM(self.v0, self.a, self.b,self.sigma,self.delta,self.s0, self.T,self.size)

            if(key in fwd):
                print("in!!")
                predictedPath = idm.getTrajectory(me,fwd[key],value[0], value[1])
                predictPath[key] = predictedPath
            else:
                print(value)
                predictedPath = idm.getTrajectory(me,None,value[0], value[1])
                predictPath[key] = predictedPath
        return predictPath
    
    def turningPositionDistribution(self,predictPath:Dict):
        """
        计算每条车道的预测位置的分布
        """
        DistributionDict = {}
        for key,predictPos in predictPath.items():
            newPos = np.array([predictPos[0][-1],predictPos[1][-1]])
            cov = np.array([[self.sigma, 0], [0, self.sigma]])
            # deltaX = 0.1
            # x,y = np.mgrid[bwd.pos[0]-deltaX:bwd.pos[0]+deltaX:.01, bwd.pos[1]-deltaX:bwd.pos[1]+deltaX:.01]  # 做测试使用的x,y，宽度定为±1
            # debug先查测试集
            
            rv = multivariate_normal(newPos, cov)
            DistributionDict[key] = rv
        return DistributionDict

def main():
    dataset = '/home/mdmp/ppp/dataset/uncompressed_scenario_training_training.tfrecord-00266-of-01000'
    dataloader = Dataloader(dataset)
    initLane = lane(0,0,0,0,0)
    meId = 161
    meIDM = IDM(14,10,4,1)
    f_new1IDM = IDM(14,10,4,1)
    me = Moveable(1, [0, 0], [0, 0, 0], 1, 4.5,4)
    me = moveableInitialize(dataloader,meId,me,meIDM)
    
    idFold = 0
    idBnew = 0
    idFnew = 0
    # 初始化当前前车
    f_old = Moveable(1, [0, 0], [0, 0, 0], initLane, 4.5,4)
    # 初始化变道后后车
    b_new1 = Moveable(1, [0, 0], [0, 0, 0], initLane, 4.5,4)
    # 初始化变道后前车
    f_new1 = Moveable(1, [0, 0], [0, 0, 0], initLane, 4.5,4)
    # 初始化向右变道后后车
    b_new2 = Moveable(1, [0, 0], [0, 0, 0], initLane, 4.5,4)
    # 初始化变道后前车
    f_new2 = Moveable(1, [0, 0], [0, 0, 0], initLane, 4.5,4)

    turning = TurningModel(0.8,14,10,4,1,1,20)
    print("###########################Initialize#####################################")
    for timeIndex in range(1,71):
        print("###############################time index:",timeIndex,"################################")
        # 车辆信息更新
        me = moveableUpdate(dataloader,meId,timeIndex,me,meIDM)
        
        print("fwd cars",dataloader.get_front_agent_inTargetLane_byID(meId,timeIndex))
        Fwd = {}
        for key,value in dataloader.get_front_agent_inTargetLane_byID(meId,timeIndex).items():
            fwd = Moveable(f_new1IDM, [0, 0], [0, 0, 0], initLane, 4.5,4)
            fwd = moveableUpdate(dataloader,value,timeIndex,fwd,f_new1IDM)
            Fwd[key] = fwd
        print(Fwd)
        pathDict = turning.turningDecision(me,dataloader.get_target_lane_points_byID(me.lane.id),Fwd)
        print('pathDict',pathDict)
        for key,value in pathDict.items():
            print("value",value)
            plot_tool.visualize_on_step(dataset,dataloader,value,timeIndex)
        me = moveableUpdate(dataloader,meId,timeIndex+1,me,meIDM)
        
        print("实际位置：",me.pos)

if __name__ == "__main__":
    main()
            
        
        



